#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
智能查询规划器模块

根据问题类型和分析轮次动态选择查询策略
"""

import logging
from typing import List, Optional
from llama_index.core import PromptTemplate
from .data_structures import AnalysisContext, QueryResult
from .enums import QueryType
from biz.core.ai.prompts import SINGLE_QUERY_PLANNING_PROMPT

logger = logging.getLogger(__name__)


class IntelligentQueryPlanner:
    """智能查询规划器 - 根据问题类型和分析轮次动态选择查询策略"""
    
    def __init__(self, llm):
        self.llm = llm
        self.single_query_planning_prompt = PromptTemplate(SINGLE_QUERY_PLANNING_PROMPT)
    
    async def plan_next_query(self, analysis_context: AnalysisContext) -> Optional[QueryType]:
        """规划下一个查询（单查询模式）"""
        try:
            # 准备上下文信息
            executed_queries = [r.query_type.value for r in analysis_context.previous_results]
            previous_summary = self._summarize_previous_results(analysis_context.previous_results)
            information_gaps = self._analyze_information_gaps(analysis_context)
            
            # 构建提示
            prompt = self.single_query_planning_prompt.format(
                user_input=analysis_context.user_input,
                problem_type=analysis_context.problem_type.value,
                device_ids=", ".join(analysis_context.device_ids) if analysis_context.device_ids else "未指定",
                executed_queries=", ".join(executed_queries) if executed_queries else "无",
                previous_summary=previous_summary,
                information_gaps=information_gaps
            )
            
            # 调用LLM
            response = await self.llm.acomplete(prompt)
            result = self._parse_planning_response(str(response))
            
            if result and "selected_query" in result:
                query_type_str = result["selected_query"]
                try:
                    query_type = QueryType(query_type_str)
                    
                    # 验证是否已执行
                    if query_type in [r.query_type for r in analysis_context.previous_results]:
                        logger.warning(f"查询类型 {query_type.value} 已执行，跳过")
                        return None
                    
                    # 更新期望洞察
                    if "expected_insights" in result:
                        analysis_context.expected_insights = result["expected_insights"]
                    
                    logger.info(f"规划下一个查询: {query_type.value}")
                    logger.info(f"选择原因: {result.get('reasoning', '未提供')}")
                    
                    return query_type
                    
                except ValueError:
                    logger.error(f"无效的查询类型: {query_type_str}")
                    return None
            
            return None
            
        except Exception as e:
            logger.error(f"查询规划失败: {str(e)}")
            return None
    
    def _summarize_previous_results(self, previous_results: List[QueryResult]) -> str:
        """总结前轮查询结果"""
        if not previous_results:
            return "无前轮查询"
        
        summaries = []
        for result in previous_results:
            status = "成功" if result.success else "失败"
            data_count = len(result.data) if result.data else 0
            summaries.append(f"{result.query_type.value}({status}, {data_count}条数据)")
        
        return "; ".join(summaries)
    
    def _analyze_information_gaps(self, analysis_context: AnalysisContext) -> str:
        """分析信息缺口"""
        executed_types = {r.query_type for r in analysis_context.previous_results}
        all_types = {QueryType.DEVICE_BASIC_INFO, QueryType.FAULT_RECORDS, QueryType.NETWORK_STATUS}
        
        missing_types = all_types - executed_types
        
        if not missing_types:
            return "所有基础查询已完成"
        
        gaps = [qt.value for qt in missing_types]
        return f"缺失信息: {', '.join(gaps)}"
    
    def _parse_planning_response(self, response: str) -> Optional[dict]:
        """解析规划响应"""
        from .json_utils import parse_json_from_text
        
        result = parse_json_from_text(
            text=response,
            default_value=None
        )
        
        if result is None:
            logger.warning("无法从响应中提取JSON")
        
        return result
